require 'config'
require 'lib'

local rulematch = ngx.re.find
local USER_AGENT_RULES = get_rule('useragent.rule')
local SERVER_NAME_LIMIT = read_serlimit_config()

function req_limit()
    local limit_req = require "resty.limit.req"

    local http_host = ngx.var.http_host
    ser_req_limit = SERVER_NAME_LIMIT["ReqLimit"][http_host]

    if ser_req_limit ~= nil then
        rate,burst = string.match(ser_req_limit,'(.*)/(.*)')
    else
	rate,burst = string.match(GlobalReqLimit,'(.*)/(.*)')
    end

    local lim, err = limit_req.new("my_limit_req_store", tonumber(rate), tonumber(burst))
    if not lim then
        ngx.log(ngx.ERR, "failed to instantiate a resty.limit.req object: ", err)
    end

    local http_host = ngx.var.http_host
    local ATTACK_URI = ngx.var.uri
    local key = http_host..get_client_ip()..ATTACK_URI
    local delay, err = lim:incoming(key, true)

    if not delay then
        if err == "rejected" then
            ngx.log(ngx.ERR, "req limit host key: ", key)
            return say_html("req")
        end
        ngx.log(ngx.ERR, "failed to limit req: ", err)
    end

    if delay >= 0.001 then
        ngx.sleep(delay)
    end

end

function count_limit()
    local limit_count = require "count"

    local http_host = ngx.var.http_host
    ser_count_limit = SERVER_NAME_LIMIT["CountLimit"][http_host]
    if ser_count_limit ~= nil then
        count,sec,interval = string.match(ser_count_limit,'(.*)/(.*)/(.*)')
    else
        count,sec,interval = string.match(GlobalCountLimit,'(.*)/(.*)/(.*)')
    end
    ngx.log(ngx.ERR, count,sec,interval)
    local lim, err = limit_count.new("my_limit_count_store", tonumber(count), tonumber(sec), tonumber(interval))
    if not lim then
        ngx.log(ngx.ERR, "failed to instantiate a resty.limit.count object: ", err)
	return ngx.exit(500)
    end
    local http_host = ngx.var.http_host
    local ATTACK_URI = ngx.var.uri
    local key = http_host..get_client_ip()..ATTACK_URI
    local delay, err = lim:incoming(key, true)

    if not delay then
        if err == "rejected" then
            ngx.log(ngx.ERR, "limit host key: ", key)
            say_html("count")
        end
        ngx.log(ngx.ERR, "failed to limit count: ", err)
    end

    local remaining = err

    ngx.header["X-RateLimit-Limit"] = "5000"
    ngx.header["X-RateLimit-Remaining"] = remaining
end

--deny user agent
function user_agent_attack_check()
    if config_user_agent_check == "on" then
        local USER_AGENT = ngx.var.http_user_agent
        if USER_AGENT ~= nil then
            for _,rule in pairs(USER_AGENT_RULES) do
                if rule ~="" and rulematch(USER_AGENT,rule,"jo") then
                    if config_waf_enable == "on" then
                        waf_output()
                        return true
                    end
                end
            end
        end
    end
    return false
end
